import cv2
import os
import shutil
import numpy as np
from PIL import Image

#  加载人脸检测器
face_cascade = cv2.CascadeClassifier('res/haarcascade_frontalface_default.xml')


def hamming_distance(str1, str2):
    if len(str1) != len(str2):
        return
    count = 0
    for i in range(len(str1)):
        if str1[i] != str2[i]:
            count += 1
    return count




def dhash(image):
    #  转为灰度图像
    image = cv2.resize(image, (8, 8), interpolation=cv2.INTER_CUBIC)
    #  缩放为更小的尺寸方便计算
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    s = 0
    hash_str = ''
    for i in range(8):
        for j in range(8):
            s = s + gray[i, j]
    avg = s / 64
    for i in range(8):
        for j in range(8):
            if gray[i, j] > avg:
                hash_str = hash_str + '1'
            else:
                hash_str = hash_str + '0'

    return hash_str


#  创建保存相似人脸的文件夹
def create_folder(path):
    if not os.path.exists(path):
        os.makedirs(path)


#  加载视频文件
video_path = 'C:/Users/17938/Documents/WXWork/1688853076112674/Cache/File/2023-07/dataset/英语.mp4'

cap = cv2.VideoCapture(video_path)

#  用于保存当前帧的人脸信息  {哈希值:  [人脸图像,  目标文件夹]}
faces_dict = {}

frame_count = 0
while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break

    #  转换为灰度图像进行人脸检测
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

    #  人脸检测
    faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(60, 60))

    for (x, y, w, h) in faces:
        #  提取人脸图像
        face_img = frame[y:y + h, x:x + w]

        #  计算图像的哈希值
        hash_value = dhash(face_img)

        #  计算图像的哈希值
        hash_value = dhash(face_img)

        if len(faces_dict) == 0:
            #  创建第一个相似人脸文件夹
            folder_path = os.path.join('D:/1/', str(hash_value))
            create_folder(folder_path)
            face_path = os.path.join(folder_path, 'face{}.jpg'.format(frame_count))
            cv2.imwrite(face_path, face_img)
            faces_dict[hash_value] = folder_path
        else:
            #  判断是否有相似人脸已经存在
            folder_paths = list(faces_dict.values())
            for folder_path in folder_paths:
                #  计算当前人脸图像与已存在文件夹中的人脸图像的汉明距离
                face_paths = os.listdir(folder_path)
                if len(face_paths) > 0:
                    first_face_path = os.path.join(folder_path, face_paths[0])
                    first_face_img = cv2.imread(first_face_path)
                    first_hash_value = dhash(first_face_img)
                    distance = hamming_distance(hash_value, first_hash_value)
                    if distance <= 20:  # 汉明距离小于等于5认为是相似人脸
                        face_path = os.path.join(folder_path, 'face{}.jpg'.format(frame_count))
                        cv2.imwrite(face_path, face_img)
                        break
            else:
                #  与所有已存在文件夹中的人脸图像都不相似，创建新的文件夹
                folder_path = os.path.join('D:/1/', str(hash_value))
                create_folder(folder_path)
                face_path = os.path.join(folder_path, 'face{}.jpg'.format(frame_count))
                cv2.imwrite(face_path, face_img)
                faces_dict[hash_value] = folder_path

        frame_count += 1

cap.release()
cv2.destroyAllWindows()